import xlrd
import numpy as np
# import scipy.io as sio
from matplotlib import pyplot as plt
import matplotlib.patches as patches
import os
import time
import collections
import itertools
import Spin_Texture_Initial_Settings

import sys
import Rb87BEC_Basic_Functions as bf
import Rb87BEC_Fit as bffit

import socket

os.system('chcp 65001')
Host_Name = socket.gethostname()


Computer_Name = 'C:\\Users\\DELL\\Desktop\\Fig2 Original Data\\'
# Index = 0.05

# %%
time_start = time.time()
time_start_struct = time.localtime(time_start)

# %% Switches
FitBEC_Switch = 1
FitDraw_Switch = 0
OriImage_Switch = 1
Data_Total = []
Fit_SaveName = ['NoFit', '']
Data_Cut_Shape_list = ['Rectangle', 'Rectangle']
# %% Modes
File_Mode = 1
Atom_Mode = 0

# %% Initial settings
'''
No need to -1 on row number.
'''
# Computer_Name = r'\\calcenter2' + '\\'

Mat_Data = np.array(range(806, 806 + 1, 1))
# Mat_Data = np.concatenate((np.array([658, 659, 683]), Mat_Data))
# Mat_Data = [683]
# Mat_Data.append(292)
Data_Exl = xlrd.open_workbook(Computer_Name + 'spin_texture.xlsx')
Table = Data_Exl.sheets()[0]
# Matrix_Shift = [0, 0]
clims = [-0.1, 0.5]
FBZ_Len = 66
Ref_Photo_Num = (20190708, 'Line47')
# Wait_Sec = 1
# Photo_Name_Tag='22860test_'
# Photo_Total = ''
# if File_Mode == 0:
#    PhotoName_Pre = 'ascOD'+Photo_Name_Tag
#    PhotoName_End = '.asc'
#    
# elif File_Mode == 1:
#    PhotoName_Pre = Photo_Name_Tag
#    PhotoName_End = 'od.mat'

# %% Calculate all photos

for Row_Num in range(len(Mat_Data)):

    Fails = []
    Data_ThisRow = []
    Exl_RowNum = Mat_Data[Row_Num] - 1
    Input_Matrix = Table.row_values(Exl_RowNum)
    Parameters = {}

    # Set parameters
    YearMonthDay = str(int(Input_Matrix[0]))
    Parameters['YearMonthDay'] = YearMonthDay
    Ini_Photo = int(Input_Matrix[1])
    Parameters['Ini_Photo'] = Ini_Photo
    Fin_Photo = int(Input_Matrix[2])
    Parameters['Fin_Photo'] = Fin_Photo
    Repeat = int(Input_Matrix[3])
    Parameters['Repeat'] = Repeat
    V0 = Input_Matrix[4]
    Parameters['V0'] = V0
    Vy = Input_Matrix[5]
    Parameters['Vy'] = Vy
    Omega = Input_Matrix[6]
    Parameters['Omega'] = Omega
    Detuning = Input_Matrix[10]
    Parameters['Detuning'] = round(Detuning, 3)
    DT6_AMP = Input_Matrix[9]
    Parameters['DT6_AMP'] = round(DT6_AMP, 3)
    Cam = Input_Matrix[8]
    if Cam == '22834':
        continue
    Photo_Name_Tag = Cam + 'test_'
    Photo_Total = ''
    if File_Mode == 0:
        PhotoName_Pre = 'ascOD' + Photo_Name_Tag
        PhotoName_End = '.asc'
    elif File_Mode == 1:
        PhotoName_Pre = Photo_Name_Tag
        PhotoName_End = 'od.mat'
    Parameters_str = str(Parameters)
    Parameters_str = Parameters_str.replace('{', '')
    Parameters_str = Parameters_str.replace('}', '')
    Parameters_str = Parameters_str.replace(' ', '')
    Parameters_str = Parameters_str.replace(':', '=')
    Photo = str(Ini_Photo) + 'To' + str(Fin_Photo)
    Photo_Total += Photo + ';'
    # Data matrices
    # if Row_Num == 0:
    #     Time = np.linspace(t_Ini, t_End, (t_End-t_Ini)/Time_Step+1)

    # Addresses
    if File_Mode == 0:
        OriData_FolderName = r'\\CAMERA-PC\ShareFolder' + '\\' + \
                                YearMonthDay + 'morning\\' + Cam + '\\'
    elif File_Mode == 1:
        OriData_FolderName = '\\\\calcenter2\\' + 'oriData' + '\\' + YearMonthDay + '\\'
    Save_Folder_Name = '\\\\calcenter2\\' + 'oriData' + '\\' + YearMonthDay + '\\'
    SaveFigure_Address = '\\\\calcenter2\\' + 'Rb87figures\\3DSOC_spin_texture\\figures' + \
                         Fit_SaveName[FitBEC_Switch] + '\\' + YearMonthDay + '\\' + Photo + '\\'
    os.system('mkdir ' + '\"' + SaveFigure_Address + '\"')
    SaveData_Address = Computer_Name + 'Rb87Data\\3DSOC_spin_texture\\' + \
        Fit_SaveName[FitBEC_Switch] + '\\' + YearMonthDay + '\\' + Photo + '\\'
    os.system('mkdir ' + '\"' + SaveData_Address + '\"')
    SaveFigure_Address_All = '\\\\calcenter2\\' + 'Rb87figures\\3DSOC_spin_texture\\figures' + \
                             Fit_SaveName[FitBEC_Switch] + '\\' + YearMonthDay + '\\' + 'All_Spin_Polarization' + '\\'
    os.system('mkdir ' + '\"' + SaveFigure_Address_All + '\"')
    SaveData_Address_All = Computer_Name + 'Rb87Data\\3DSOC_spin_texture\\' + \
        Fit_SaveName[FitBEC_Switch] + '\\' + YearMonthDay + '\\' + 'All_Spin_Polarization' + '\\'
    os.system('mkdir ' + '\"' + SaveData_Address_All + '\"')
    Atom_Center = Spin_Texture_Initial_Settings.AtomCenter_Pos()
    Cut_Width = Spin_Texture_Initial_Settings.CutWidths()

    # Calculating
    Photo_Ind = Ini_Photo
    SPs = collections.OrderedDict()
    FBZ_Ups = collections.OrderedDict()
    FBZ_Dns = collections.OrderedDict()
    xs = collections.OrderedDict()
    while Photo_Ind <= Fin_Photo:
        plt.close('all')
        Load_Name = OriData_FolderName + PhotoName_Pre + '%03d' % Photo_Ind + PhotoName_End
        Save_Name = Save_Folder_Name + PhotoName_Pre + '%03d' % Photo_Ind + 'od'
        #        if Cam == '22860':
        #            All_x = [0, 1024]
        #            All_y = [0, 1024]
        #            Up_x = [458, 658]
        #            Up_y = [480, 680]
        #            Dn_x = [517, 717]
        #            Dn_y = [230, 430]
        #        elif Cam == '4552':
        #            All_x = [0, 512]
        #            All_y = [0, 512]
        #            Up_x = [150, 330]
        #            Up_y = [320, 500]
        #            Dn_x = [240, 400]
        #            Dn_y = [170, 330]
        #        else:
        #            break
        Data_ThisPic_AtomNum, Data_ThisPic_AtomRegion, ni = \
            bf.Load_Pic_And_Cal(Photo_Ind, Load_Name, File_Mode, Save_Name, SaveFigure_Address + 'All_' + Cam + '_',
                                Atom_Center, Cut_Width, len(Cut_Width) - 1, OriImage_Switch,
                                Rot_Ang=43, Enclose=Data_Cut_Shape_list)
        Fit_Regions = []
        Fit_AtomNums = []
        TPx = [[], []]
        Sum_0 = np.sum(Data_ThisPic_AtomRegion[0])
        Sum_1 = np.sum(Data_ThisPic_AtomRegion[1])
        Discard_Pic_Threshold = 250
        #        FitBEC_Switch = 1
        #        if Sum_0 < Discard_Pic_Threshold or Sum_1 < Discard_Pic_Threshold:
        #            print('Sum = %.2f and %.2f, no need to fit'%(Sum_0, Sum_1))
        #            FitBEC_Switch = 0
        if FitBEC_Switch == 1:
            for Atom_Ind in range(2):
                AtomRegion = Data_ThisPic_AtomRegion[Atom_Ind].copy()

                plt.imshow(AtomRegion, vmin=-0.1, vmax=0.5, cmap='jet')
                plt.title('%d, %d' % (Photo_Ind, Atom_Ind))
                plt.pause(0.01)
                IsBEC_Threshold = [0.010, 0.015]
                Region_Size = np.shape(AtomRegion)
                AtomCloud_Size = bf.Cal_AtomSize(AtomRegion, [50, 0.20])
                if AtomCloud_Size == 0:
                    AtomCloud_Size = 10000
                Od_Max = np.max(AtomRegion)
                Fi = Od_Max / AtomCloud_Size
                BEC_Judge = bf.Is_BEC(AtomRegion, AtomCloud_Size, IsBEC_Threshold)
                print('BEC_Judge =', BEC_Judge)
                if np.sum(AtomRegion) < Discard_Pic_Threshold and BEC_Judge == 0:
                    print("The atom cloud is too hot. Quit fitting.")
                    Fit_Region = AtomRegion.copy()
                    Fit_Regions.append(Fit_Region)
                    Fit_AtomNums.append(np.sum(Fit_Region) * 112.2)
                    Check_x = np.array(
                        [[1, 0, 0, 100, 100, 0, 3, 0, 30, 30, 0.5], [1, 0, 0, 100, 100, 0, 3, 0, 30, 30, 0.5],
                         [1, 0, 0, 100, 100, 0, 3, 0, 30, 30, 0.5]])
                    Check_x = Check_x.T
                    TPx.append(Check_x)
                    continue
                Fitting_Parameters = np.zeros((11, 3))
                if BEC_Judge == 1:
                    Od_Lim = 0.3
                    Fitting_Parameters[0, :] = [1, 0, 2]  # Amp
                    Fitting_Parameters[1, :] = [0, -Region_Size[0] / 2, Region_Size[0] / 2]  # x0
                    Fitting_Parameters[2, :] = [0, -Region_Size[1] / 2, Region_Size[1] / 2]  # y0
                    Fitting_Parameters[3, :] = [100, 10, Region_Size[0]]  # wx
                    Fitting_Parameters[4, :] = [100, 10, Region_Size[1]]  # wy
                    Fitting_Parameters[5, :] = [0, -np.pi, np.pi]  # rot_ang
                    Fitting_Parameters[6, :] = [3, 0, 10]  # A2
                    Fitting_Parameters[7, :] = [0, -1, 1]  # offset
                    Fitting_Parameters[8, :] = [10, 0, 20]  # wxc
                    Fitting_Parameters[9, :] = [10, 0, 20]  # wyc
                    Fitting_Parameters[10, :] = [0.5, 0, 1]  # Ac
                    x_Start = Fitting_Parameters[:, 0]
                    x_LowBound = Fitting_Parameters[:, 1]
                    x_UpBound = Fitting_Parameters[:, 2]
                elif BEC_Judge == 0.5:
                    Od_Low = 0.5
                    Od_Lim = np.max((np.max(AtomRegion * 0.5), Od_Low))
                    Size_Start = np.max((AtomCloud_Size, 11))
                    Fitting_Parameters[0, :] = [Od_Lim / 2, 0, Od_Lim]  # Amp
                    Fitting_Parameters[1, :] = [0, -Region_Size[0] / 2, Region_Size[0] / 2]  # x0
                    Fitting_Parameters[2, :] = [0, -Region_Size[1] / 2, Region_Size[1] / 2]  # y0
                    Fitting_Parameters[3, :] = [Size_Start, 10, Region_Size[0]]  # wx
                    Fitting_Parameters[4, :] = [Size_Start, 10, Region_Size[1]]  # wy
                    Fitting_Parameters[5, :] = [0, -np.pi, np.pi]  # rot_ang
                    Fitting_Parameters[6, :] = [Od_Lim / 2, 0, Od_Lim]  # A2
                    Fitting_Parameters[7, :] = [0, -1, 1]  # offset
                    Fitting_Parameters[8, :] = [30, 10, 50]  # wxc
                    Fitting_Parameters[9, :] = [30, 10, 50]  # wyc
                    Fitting_Parameters[10, :] = [Od_Lim / 2, 0, Od_Lim]  # Ac
                    x_Start = Fitting_Parameters[:, 0]
                    x_LowBound = Fitting_Parameters[:, 1]
                    x_UpBound = Fitting_Parameters[:, 2]
                else:
                    Od_Lim = 1
                    Fitting_Parameters[0, :] = [Od_Lim / 2, 0, Od_Lim]  # Amp
                    Fitting_Parameters[1, :] = [0, -Region_Size[0] / 2, Region_Size[0] / 2]  # x0
                    Fitting_Parameters[2, :] = [0, -Region_Size[1] / 2, Region_Size[1] / 2]  # y0
                    Fitting_Parameters[3, :] = [AtomCloud_Size, 0, Region_Size[0]]  # wx
                    Fitting_Parameters[4, :] = [AtomCloud_Size, 0, Region_Size[1]]  # wy
                    Fitting_Parameters[5, :] = [0, -np.pi, np.pi]  # rot_ang
                    Fitting_Parameters[6, :] = [Od_Lim / 2, 0, Od_Lim]  # A2
                    Fitting_Parameters[7, :] = [0, -1, 1]  # offset
                    Fitting_Parameters[8, :] = [7, 0, 40]  # wxc
                    Fitting_Parameters[9, :] = [7, 0, 40]  # wyc
                    Fitting_Parameters[10, :] = [Od_Lim / 2, 0, Od_Lim]  # Ac
                    x_Start = Fitting_Parameters[:, 0]
                    x_LowBound = Fitting_Parameters[:, 1]
                    x_UpBound = Fitting_Parameters[:, 2]
                try:
                    Fit_Region, Check_x = bffit.Main_Gauss(AtomRegion, x_Start, x_LowBound, x_UpBound, Od_Lim,
                                                           'Bimodal')
                except:
                    Fails.append(Photo_Ind)
                    print('Fit failed. Number saved. Check again later.\n')
                    Fit_Region = np.zeros(Region_Size)
                Fit_Regions.append(Fit_Region)
                Fit_AtomNums.append(np.sum(Fit_Region) * 112.2)
                TPx.append(Check_x)
            Fit_AtomNumsArray = np.array(Fit_AtomNums)
            Data_ThisPic_AtomNum = Fit_AtomNumsArray / np.sum(Fit_AtomNumsArray)
            Data_ThisPic_AtomNum = Data_ThisPic_AtomNum.tolist()
            Data_ThisPic_AtomRegion = Fit_Regions.copy()

        Data_Up = Data_ThisPic_AtomRegion[0].copy()
        Data_Dn = Data_ThisPic_AtomRegion[1].copy()
        Img_Len_Up = Cut_Width[0]
        Img_Len_Dn = Cut_Width[1]

        DUL_N = 4
        DUL_Range = DUL_N * FBZ_Len
        Data_Up_Large = np.zeros((DUL_Range, DUL_Range))
        Data_Up_Large[(DUL_Range - Img_Len_Up) // 2:(DUL_Range + Img_Len_Up) // 2,
                      (DUL_Range - Img_Len_Up) // 2:(DUL_Range + Img_Len_Up) // 2] = Data_Up
        FBZ_Starts_Up = list(itertools.product(np.arange(DUL_N) * FBZ_Len, repeat=2))

        BZs_Up = []
        fig, ax = plt.subplots(1)
        ax.imshow(Data_Up_Large, vmin=-0.1, vmax=0.5, cmap='jet')
        for Starts_Up in FBZ_Starts_Up:
            Ends_Up = tuple(np.array(Starts_Up) + FBZ_Len)
            BZs_Up.append(Data_Up_Large[Starts_Up[0]:Ends_Up[0], Starts_Up[1]:Ends_Up[1]])
            rect = patches.Rectangle((Starts_Up[0], Starts_Up[1]), FBZ_Len, FBZ_Len, linewidth=1, edgecolor='r',
                                     facecolor='none')
            ax.add_patch(rect)
        BZs_Up = np.array(BZs_Up)
        FBZ_Up = np.sum(BZs_Up, 0)
        plt.show()
        fig.savefig(SaveFigure_Address + 'Cam_' + Cam + '_' + str(Photo_Ind) + r'Spin_Up' + '.png')
        fig.savefig(SaveFigure_Address+'Cam_'+Cam+'_'+str(Photo_Ind)+r'Spin_Up'+'.pdf')

        DDL_N = 3
        DDL_Range = DDL_N * FBZ_Len
        Data_Dn_Large = np.zeros((DDL_Range, DDL_Range))
        Data_Dn_Large[(DDL_Range - Img_Len_Dn) // 2:(DDL_Range + Img_Len_Dn) // 2,
                      (DDL_Range - Img_Len_Dn) // 2:(DDL_Range + Img_Len_Dn) // 2] = Data_Dn
        FBZ_Starts_Dn = list(itertools.product(np.arange(DDL_N) * FBZ_Len, repeat=2))
        BZs_Dn = []
        fig, ax = plt.subplots(1)
        ax.imshow(Data_Dn_Large, vmin=-0.1, vmax=0.5, cmap='jet')
        for Starts_Dn in FBZ_Starts_Dn:
            Ends_Dn = tuple(np.array(Starts_Dn) + FBZ_Len)
            BZs_Dn.append(Data_Dn_Large[Starts_Dn[0]:Ends_Dn[0], Starts_Dn[1]:Ends_Dn[1]])
            rect = patches.Rectangle((Starts_Dn[0], Starts_Dn[1]), FBZ_Len, FBZ_Len, linewidth=1, edgecolor='r',
                                     facecolor='none')
            ax.add_patch(rect)
        BZs_Dn = np.array(BZs_Dn)
        FBZ_Dn = np.sum(BZs_Dn, 0)
        plt.show()
        fig.savefig(SaveFigure_Address + 'Cam_' + Cam + '_' + str(Photo_Ind) + r'Spin_Down' + '.png')
        fig.savefig(SaveFigure_Address+'Cam_'+Cam+'_'+str(Photo_Ind)+r'Spin_Down'+'.pdf')

        fig, ax = plt.subplots(1)
        plt.imshow(FBZ_Up, vmin=-0.1, vmax=0.5, cmap='seismic')
        plt.colorbar()
        plt.title('Spin Up')
        plt.show()
        fig.savefig(SaveFigure_Address + 'Cam_' + Cam + '_' + str(Photo_Ind) + r'Spin_Up_FBZ' + '.png')
        fig.savefig(SaveFigure_Address+'Cam_'+Cam+'_'+str(Photo_Ind)+r'Spin_Up_FBZ'+'.pdf')

        fig, ax = plt.subplots(1)
        plt.imshow(FBZ_Dn, vmin=-0.1, vmax=0.5, cmap='seismic')
        plt.colorbar()
        plt.title('Spin Down')
        plt.show()
        fig.savefig(SaveFigure_Address + 'Cam_' + Cam + '_' + str(Photo_Ind) + r'Spin_Down_FBZ' + '.png')
        fig.savefig(SaveFigure_Address+'Cam_'+Cam+'_'+str(Photo_Ind)+r'Spin_Down_FBZ'+'.pdf')

        fig, ax = plt.subplots(1)
        #        SPs.append(((FBZ_Up-FBZ_Dn)/(FBZ_Up+FBZ_Dn), Photo_Ind))
        SPs['%03d' % Photo_Ind] = (FBZ_Up - FBZ_Dn) / (FBZ_Up + FBZ_Dn)
        FBZ_Ups['%03d' % Photo_Ind] = FBZ_Up
        FBZ_Dns['%03d' % Photo_Ind] = FBZ_Dn
        xs['%03d' % Photo_Ind] = TPx

        np.savez(SaveData_Address_All + 'Cam_' + Cam + '_photo' + str(Photo_Ind) + r'_Spin_Polarization_FBZ',
                 SP=(FBZ_Up - FBZ_Dn) / (FBZ_Up + FBZ_Dn), FBZ_Up=FBZ_Up, FBZ_Dn=FBZ_Dn,
                 Check_xs_up=xs['%03d' % Photo_Ind][0], Check_xs_dn=xs['%03d' % Photo_Ind][1])

        plt.imshow((FBZ_Up - FBZ_Dn) / (FBZ_Up + FBZ_Dn), vmin=-1, vmax=1, cmap='seismic')
        plt.colorbar()
        plt.title('Spin Polarization')
        plt.show()
        fig.savefig(
            SaveFigure_Address + 'Cam_' + Cam + '_' + str(Photo_Ind) + r'Spin_Polarization_FBZ_color-1to1' + '.png')
        fig.savefig(SaveFigure_Address+'Cam_'+Cam+'_'+str(Photo_Ind)+r'Spin_Polarization_FBZ'+'.pdf')

        fig, ax = plt.subplots(1)
        plt.imshow((FBZ_Up - FBZ_Dn) / (FBZ_Up + FBZ_Dn), vmin=-1, vmax=1, cmap='seismic')
        plt.colorbar()
        plt.axis('off')
        plt.title('detuning=%.2f' % Parameters['Detuning'])
        plt.show()
        fig.savefig(
            SaveFigure_Address + 'Cam_' + Cam + '_' + str(Photo_Ind) + r'Spin_Polarization_FBZ_color-1to1_det' + '.png')
        fig.savefig(SaveFigure_Address+'Cam_'+Cam+'_'+str(Photo_Ind)+r'Spin_Polarization_FBZ_color-1to1_det'+'.pdf')
        Photo_Ind += 1
    plt.close('all')
    ii = 1

    fig = plt.figure(Row_Num, figsize=(20, 10))
    #    plt.figure(figsize=(20, 10))
    for SPkey in SPs.keys():
        SP = SPs[SPkey]
        plt.subplot((Fin_Photo - Ini_Photo + 1) // 10 + 1, 10, ii)
        plt.imshow(SP, vmin=-1, vmax=1, cmap='seismic')
        plt.axis('off')
        plt.title('p%s' % SPkey, fontsize=16)
        plt.tight_layout()
#        plt.subplots_adjust(top=0.9, bottom=0.1, left=0.021, right=0.992, hspace=0.0, wspace=0.119)
        ii += 1
    # plt.suptitle(Parameters_str, fontsize=20)
    plt.suptitle("Row=%d, Detuning=%.3f" % (Mat_Data[Row_Num], Parameters['Detuning']), fontsize=20)
    plt.show()
    fig.savefig(SaveFigure_Address + 'Cam_' + Cam + '_' + Photo + r'Spin_Polarization_FBZ' + '.png')
    fig.savefig(SaveFigure_Address+'Cam_'+Cam+'_'+Photo+r'Spin_Polarization_FBZ'+'.pdf')
    np.savez(SaveData_Address + 'Cam_' + Cam + '_' + Photo + r'Spin_Polarization_FBZ', SPs=SPs,
             FBZ_Ups=FBZ_Ups, FBZ_Dns=FBZ_Dns, xs=xs, Fails=Fails)

    fig.savefig(SaveFigure_Address_All + 'Cam_' + Cam + '_' + Photo + 'Row=%d' % Mat_Data[
        Row_Num] + r'Spin_Polarization_FBZ' + '.png')
    fig.savefig(SaveFigure_Address_All + 'Cam_' + Cam + '_' + Photo + 'Row=%d' % Mat_Data[
        Row_Num] + r'Spin_Polarization_FBZ' + '.pdf')
    np.savez(SaveData_Address_All + 'Cam_' + Cam + '_' + Photo + r'Spin_Polarization_FBZ', SPs=SPs,
             FBZ_Ups=FBZ_Ups, FBZ_Dns=FBZ_Dns, xs=xs,
             Fails=Fails)
time_end = time.time()
# np.savez('D:\\Rb87figures\\3DSOC_spin_texture\\figures\\20190430Fails154', Fails=Fails)
print('Time spent: ', time_end - time_start, 's')
time_end_struct = time.localtime(time_end)
print('Start time:', time_start_struct)
print('End time:', time_end_struct)
